//
// Created by hujin on 2023/3/9.
//

#ifndef SI_NODE_STATS_H
#define SI_NODE_STATS_H


#include <utility>

template <typename T> struct NodeStat{
    T volume;
    T p_o; //Probability of RW from outside of the module to inside of the module.
    NodeStat(T volume, T pO) : volume(volume), p_o(pO) {}
};

template <typename T> struct NodeStatL3 : NodeStat<T>{
    T volume_h;
    T p_a; //Probability of RW from one submodule to another submodule.
    NodeStatL3(T volume, T pO, T volumeH, T pA) : NodeStat<T>(volume, pO), volume_h(volumeH), p_a(pA) {}
};

template <typename T> struct NodeStatL2 : NodeStat<T>{
    T volume_h;
public:
    explicit operator NodeStatL3<T>() const{
        return {NodeStat<T>::volume, NodeStat<T>::p_o, volume_h, NodeStat<T>::volume};
    }

    NodeStatL2(T volume, T pO, T volumeH) : NodeStat<T>(volume, pO), volume_h(volumeH) {}
};

template <typename T> class Update{
    inline static NodeStatL3<T> push(NodeStatL3<T> &t);
    inline static NodeStatL3<T> push(NodeStatL2<T> &t);
    inline static NodeStatL3<T> merge(NodeStatL3<T> &n1, NodeStatL3<T> &n2, T e12, NodeStatL3<T> *n_h = nullptr);
    inline static NodeStatL2<T> merge(NodeStatL2<T> &n1, NodeStatL2<T> &n2, T e12, NodeStatL3<T> *n_h = nullptr);
    inline static void swap(NodeStatL3<T> &s, NodeStatL3<T> &t,  NodeStatL3<T> &child, T esc, T ect, NodeStatL3<T> *n_h = nullptr);
    inline static void swap(NodeStatL2<T> &s, NodeStatL2<T> &t,  NodeStat<T> &leaf, T esc, T ect, NodeStatL3<T> *n_h = nullptr);
};

template<typename T>
void Update<T>::swap(NodeStatL2<T> &s, NodeStatL2<T> &t, NodeStat<T> &leaf, T esc, T ect, NodeStatL3<T> *n_h) {
    t.volume += leaf.volume;
    s.volume -= leaf.volume;

    auto p0 = leaf.p_o;
    t.p_o += p0 - ect;
    s.p_o -= p0 - esc;

    if(n_h != nullptr) n_h->p_a += esc - ect;
}


template<typename T>
void Update<T>::swap(NodeStatL3<T> &s, NodeStatL3<T> &t, NodeStatL3<T> &child, T esc, T ect,  NodeStatL3<T> *n_h ) {
    t.volume += child.volume;
    s.volume -= child.volume;
    t.p_a -= child.p_o;
    s.p_a += child.p_o;

    auto p0 = child.p_o;
    t.p_o += p0 - ect;
    s.p_o -= p0 - esc;

    if(n_h != nullptr) n_h->p_a += esc - ect;
}


template<typename T>
NodeStatL2<T> Update<T>::merge(NodeStatL2<T> &n1, NodeStatL2<T> &n2, T e12, NodeStatL3<T> *n_h) {
    if(n_h != nullptr) n_h->p_a -= e12;
    //todo : computation correction
    return NodeStatL2<T>{n1.volume + n2.volume, n1.p_o + n2.p_o - e12, n1.volume_h};
}

template<typename T>
NodeStatL3<T> Update<T>::merge(NodeStatL3<T> &n1, NodeStatL3<T> &n2, T e12, NodeStatL3<T> *n_h) {
    if(n_h != nullptr) n_h->p_a -= e12;
    //todo : numerical correction
    return NodeStatL3<T>{n1.volume + n2.volume, n1.p_o + n2.p_o - e12, n1.volume_h, n1.p_a + n2.p_a};
}

template<typename T>
NodeStatL3<T> Update<T>::push(NodeStatL2<T> &t) {
    auto vh = t.volume_h;
    t.volume_h = t.volume;
    return NodeStatL3<T>{t.volume, t.p_o, vh, t.volume};
}

template<typename T>
NodeStatL3<T> Update<T>::push(NodeStatL3<T> &t) {
    auto vh = t.volume_h;
    t.volume_h = t.volume;
    return NodeStatL3<T>{t.volume, t.p_o, vh, t.p_a};
}







#endif //SI_NODE_STATS_H
